import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, random_split
import pickle
import numpy as np
from sklearn.metrics import accuracy_score, f1_score

# 检查GPU可用性
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"使用设备: {device}")

# 1. 加载数据集
with open('steel_defects_data.pkl', 'rb') as f:
    data = pickle.load(f)
    features_tensor = data['features']
    labels_tensor = data['labels']

    # 如果 labels 是独热编码，转换为类别索引
    if len(labels_tensor.shape) > 1 and labels_tensor.shape[1] > 1:
        labels_tensor = torch.argmax(labels_tensor, dim=1)


class SteelDefectsDataset(Dataset):
    def __init__(self, features, labels):
        self.features = features
        self.labels = labels

    def __len__(self):
        return len(self.features)

    def __getitem__(self, idx):
        return self.features[idx], self.labels[idx]


dataset = SteelDefectsDataset(features_tensor, labels_tensor)

# 2. 划分训练集和测试集 (80%训练, 20%测试)
total_size = len(dataset)
train_size = int(0.8 * total_size)
test_size = total_size - train_size

train_dataset, test_dataset = random_split(
    dataset, [train_size, test_size],
    generator=torch.Generator().manual_seed(42)  # 设置随机种子确保可复现性
)

# 3. 创建数据加载器
batch_size = 64
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

print(f"训练样本数: {len(train_dataset)}, 测试样本数: {len(test_dataset)}")
print(f"特征维度: {dataset[0][0].shape}")


# 4. 定义全连接神经网络模型
class SteelDefectsFCN(nn.Module):
    def __init__(self, input_dim, num_classes):
        super(SteelDefectsFCN, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(input_dim, 256),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(128, num_classes)
        )

    def forward(self, x):
        return self.fc(x)

    # 初始化模型


input_dim = train_dataset[0][0].shape[0]  # 获取特征维度
num_classes = len(torch.unique(labels_tensor))  # 获取类别数量
model = SteelDefectsFCN(input_dim, num_classes).to(device)

# 5. 设置训练参数
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-5)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=3, factor=0.5)


# 6. 训练函数
def train(model, loader, optimizer, criterion):
    model.train()
    total_loss = 0
    for features, labels in loader:
        features = features.float().to(device)
        labels = labels.long().to(device)

        optimizer.zero_grad()
        outputs = model(features)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        total_loss += loss.item() * features.size(0)
    return total_loss / len(loader.dataset)


# 7. 评估函数
def evaluate(model, loader):
    model.eval()
    all_preds = []
    all_labels = []
    total_loss = 0

    with torch.no_grad():
        for features, labels in loader:
            features = features.float().to(device)
            labels = labels.long().to(device)

            outputs = model(features)
            loss = criterion(outputs, labels)
            total_loss += loss.item() * features.size(0)

            _, preds = torch.max(outputs, 1)
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())

    accuracy = accuracy_score(all_labels, all_preds)
    f1 = f1_score(all_labels, all_preds, average='macro')
    return total_loss / len(loader.dataset), accuracy, f1


# 8. 训练循环
num_epochs = 1000
best_f1 = 0.0

for epoch in range(num_epochs):
    train_loss = train(model, train_loader, optimizer, criterion)
    val_loss, val_acc, val_f1 = evaluate(model, test_loader)
    scheduler.step(val_loss)

    # 保存最佳模型
    if val_f1 > best_f1:
        best_f1 = val_f1
        torch.save(model.state_dict(), 'best_model.pth')

    print(f"Epoch {epoch + 1}/{num_epochs} | "
          f"Train Loss: {train_loss:.4f} | "
          f"Val Loss: {val_loss:.4f} | "
          f"Val Acc: {val_acc:.4f} | "
          f"Val F1: {val_f1:.4f} | "
          f"LR: {optimizer.param_groups[0]['lr']:.6f}")

# 9. 加载最佳模型并最终评估
model.load_state_dict(torch.load('best_model.pth'))
test_loss, test_acc, test_f1 = evaluate(model, test_loader)
print(f"\n最终测试结果: Loss: {test_loss:.4f}, Acc: {test_acc:.4f}, F1: {test_f1:.4f}")